export DEBUG_MODE="true" # Enable Debug if you want to see the rollout of model during RL
export LOG_PATH="./debug_log_2b.txt"

EXP_NAME='500m-r1-count-grpo-clevr'
MODEL_PATH='checkpoints/namo-500m-aimv2-native-pixelshuffle_2x-sft-final-whole'

torchrun --nproc_per_node="8" \
    --nnodes="1" \
    --node_rank="0" \
    --master_addr="127.0.0.1" \
    --master_port="12345" \
    train_grpo.py \
    --output_dir ./checkpoints/namo-$EXP_NAME \
    --pretrain_model_path $MODEL_PATH \
    --dataset_name 'clevr_cogen_a_train' \
    --max_prompt_length 1024 \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 2 \
    --logging_steps 1 \
    --bf16 \
    --report_to wandb \
    --gradient_checkpointing false \
    --attn_implementation flash_attention_2 \
    --max_pixels 401408 \
    --num_train_epochs 2 \
    --run_name $EXP_NAME \
    --save_steps 100 \
    --save_only_model true